{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.indexes import GraphIndexCreator\n",
    "from langchain.llms import OpenAI\n",
    "from langchain.document_loaders import TextLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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\n",
      "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\n"
     ]
    }
   ],
   "source": [
    "# openai.api_key = 'sk-5vOFyAXiUcLGkRKV60D2D7E6Ee6e408f8d366f4c1013C004'\n",
    "# # 更换OpenAI接口的host\n",
    "# openai.api_base = \"https://api.akm.pw/v1\" #在这里设置即可,需要特别注意这里的/v1是必须的，否则报错。前面的地址注意替换即可。\n",
    "\n",
    "import os \n",
    "from dotenv import load_dotenv, find_dotenv, dotenv_values \n",
    "# os.environ[\"OPENAI_API_KEY\"] = \"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\"\n",
    "os.environ['HTTPS_PROXY']    = 'http://127.0.0.1:7890'\n",
    "os.environ[\"HTTP_PROXY\"]     = 'http://127.0.0.1:7890'\n",
    "\n",
    "# 读取本地的环境变量 \n",
    "env_vars = dotenv_values('.env')\n",
    "# 获取环境变量 OPENAI_API_KEY\n",
    "openai_api_key = env_vars['OPENAI_API_KEY']\n",
    "print(openai_api_key)\n",
    "\n",
    "# 获取环境变量的值\n",
    "api_key = os.environ.get('OPENAI_API_KEY')\n",
    "\n",
    "# 打印环境变量的值\n",
    "print(api_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "f:\\Anaconda\\envs\\torch_gpu\\lib\\site-packages\\langchain\\llms\\openai.py:173: UserWarning: You are trying to use a chat model. This way of initializing it is no longer supported. Instead, please use: `from langchain.chat_models import ChatOpenAI`\n",
      "  warnings.warn(\n",
      "f:\\Anaconda\\envs\\torch_gpu\\lib\\site-packages\\langchain\\llms\\openai.py:751: UserWarning: You are trying to use a chat model. This way of initializing it is no longer supported. Instead, please use: `from langchain.chat_models import ChatOpenAI`\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# llm = OpenAI(model_name=\"gpt-3.5-turbo\", temperature=0,\n",
    "#              openai_api_key='sk-5vOFyAXiUcLGkRKV60D2D7E6Ee6e408f8d366f4c1013C004',\n",
    "#              openai_api_base='https://api.akm.pw/v1')\n",
    "\n",
    "llm = OpenAI(model_name=\"gpt-3.5-turbo\", temperature=0,\n",
    "             openai_api_key=openai_api_key)\n",
    "\n",
    "print(llm.openai_api_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from getpass import getpass\n",
    "\n",
    "index_creator = GraphIndexCreator(llm=llm)\n",
    "\n",
    "# openai_api_key = getpass(\"Enter your OpenAI API key: \")\n",
    "\n",
    "# index_creator = IndexCreator(openai_api_key=openai_api_key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('../generation_dataset/raw/甄嬛传剧本01-10.txt') as f:\n",
    "    all_text = f.read()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['华妃：你是本宫的家生奴才，竟这般不懂规矩，本宫也不便教你了，若是皇后不饶恕你，本宫也不会轻放了你。',\n",
       " '皇后：妹妹，只是小事，不用动这么大的气呀。',\n",
       " '华妃：她本是粗笨，不机灵，幸得娘娘体恤，臣妾回去一定会好好教导她的。',\n",
       " '皇后：颂芝原是妹妹的陪嫁丫鬟，身份不同一些，怎能让她这样端茶倒水的？你若觉得颂芝不好，也不必生气，对吧？福子！',\n",
       " '福子（入）：皇后娘娘金安。华妃娘娘金安。',\n",
       " '皇后：内务府新挑来的丫头，叫福子，本宫看她机灵，便拨给你使唤吧。',\n",
       " '华妃：颂芝虽粗笨，但是翊坤宫还不缺宫女，还是皇后自己留着用吧。',\n",
       " '皇后：早听说翊坤宫的宫女做事利索，是该让福子她们这些小丫头学学了，有妹妹调教着，帮着颂芝做些粗活，也能叫她们学得乖一些。',\n",
       " '华妃：臣妾先告退了。——还不快走？！',\n",
       " '颂芝：是。']"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_text.split('\\n')[80:90]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'华妃：你是本宫的家生奴才，竟这般不懂规矩，本宫也不便教你了，若是皇后不饶恕你，本宫也不会轻放了你。\\n皇后：妹妹，只是小事，不用动这么大的气呀。\\n华妃：她本是粗笨，不机灵，幸得娘娘体恤，臣妾回去一定会好好教导她的。\\n皇后：颂芝原是妹妹的陪嫁丫鬟，身份不同一些，怎能让她这样端茶倒水的？你若觉得颂芝不好，也不必生气，对吧？福子！\\n福子（入）：皇后娘娘金安。华妃娘娘金安。\\n皇后：内务府新挑来的丫头，叫福子，本宫看她机灵，便拨给你使唤吧。\\n华妃：颂芝虽粗笨，但是翊坤宫还不缺宫女，还是皇后自己留着用吧。\\n皇后：早听说翊坤宫的宫女做事利索，是该让福子她们这些小丫头学学了，有妹妹调教着，帮着颂芝做些粗活，也能叫她们学得乖一些。\\n华妃：臣妾先告退了。——还不快走？！\\n颂芝：是。'"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "text = \"\\n\".join(all_text.split('\\n')[80:90])\n",
    "text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = index_creator.from_text(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('华妃', '本宫的家生奴才', '是'),\n",
       " ('华妃', '规矩', '不懂'),\n",
       " ('华妃', '粗笨', '是'),\n",
       " ('华妃', '娘娘', '金安'),\n",
       " ('本宫', '教', '不便'),\n",
       " ('本宫', '轻放', '不会'),\n",
       " ('皇后', '你', '不饶恕'),\n",
       " ('皇后', '妹妹', '是'),\n",
       " ('皇后', '妹妹的陪嫁丫鬟', '是'),\n",
       " ('皇后', '用', '留着'),\n",
       " ('妹妹', '着', '调教'),\n",
       " ('小事', '动这么大的气', '不用'),\n",
       " ('娘娘', '臣妾', '体恤'),\n",
       " ('臣妾', '教导', '回去'),\n",
       " ('身份', '一些', '不同'),\n",
       " ('让', '端茶倒水', '她'),\n",
       " ('让', '学学', '福子她们'),\n",
       " ('觉得', '颂芝', '不好'),\n",
       " ('颂芝', '粗笨', '虽'),\n",
       " ('生气', '对吧', '不必'),\n",
       " ('福子', '皇后娘娘金安', '入'),\n",
       " ('内务府', '丫头', '挑来'),\n",
       " ('叫', '使唤', '福子'),\n",
       " ('叫', '学得乖一些', '她们'),\n",
       " ('翊坤宫', '宫女', '不缺'),\n",
       " ('听说', '做事利索', '翊坤宫的宫女'),\n",
       " ('帮着', '做些粗活', '颂芝'),\n",
       " ('是', '颂芝', 'None')]"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph.get_triples()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Querying the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new GraphQAChain chain...\u001b[0m\n",
      "Entities Extracted:\n",
      "\u001b[32;1m\u001b[1;3m NONE\u001b[0m\n",
      "Full Context:\n",
      "\u001b[32;1m\u001b[1;3m\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' 颂芝是中国古代诗人李白的妹妹。'"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.chains import GraphQAChain\n",
    "\n",
    "chain = GraphQAChain.from_llm(OpenAI(temperature=0), graph=graph, verbose=True)\n",
    "chain.run(\"颂芝是谁？\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['皇后 不饶恕 你', '皇后 是 妹妹', '皇后 是 妹妹的陪嫁丫鬟', '皇后 留着 用']"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph.get_entity_knowledge('皇后')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch_gpu",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.15"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
